Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[TKW] Add support for multiple/local reduceOp #234

Merged
merged 3 commits into from
Oct 24, 2024

Conversation

raikonenfnu
Copy link
Contributor

@raikonenfnu raikonenfnu commented Oct 21, 2024

In order to support flash attention, we'd need to be able to expand ReduceOps in the reduction dimension as well. We will do this by expanding the source of ReduceOp and locally reduce all of them. In that effort, we introduce this PR(1st out of 2) that add support of locally reducing over multiple variables.

The second PR on the way would be expansion of ReduceOp.

In this PR we are contributing two things:

  1. Checks for consistency of indexing_dims, types, thread_shapes for multiple sources of ReduceOp
  2. Modify emit of local reduction to generate iteratively slice and reduce over multiple arguments/srcs.

Copy link
Contributor

@harsh-nod harsh-nod left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm! thanks!

Signed-off-by: Stanley Winata <[email protected]>
@harsh-nod harsh-nod merged commit 50e17a5 into iree-org:main Oct 24, 2024
7 of 8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants